Skip to content

Adapt inspect_checkpoint.py for Automated Ckpt Validation Pipeline (Agent Integration)#4567

Open
olufiyin19 wants to merge 2 commits into
mainfrom
feat/legacy-inspect-checkpoint-nnx-support
Open

Adapt inspect_checkpoint.py for Automated Ckpt Validation Pipeline (Agent Integration)#4567
olufiyin19 wants to merge 2 commits into
mainfrom
feat/legacy-inspect-checkpoint-nnx-support

Conversation

@olufiyin19

@olufiyin19 olufiyin19 commented Jul 22, 2026 •
edited
Loading

Copy link
Copy Markdown
Collaborator

Description

This PR modifies inspect_checkpoint.py to introduce native NNX architecture support, while preserving Linen compatibility as a fallback. The previous core inspect_checkpoint.py utility only supported Linen-based models. As our automated validation pipeline transitions to an NNX-first architecture, a way to inspect checkpoints against NNX abstract models is required.

  • Specific implementation details:
  1. Integrated create_nnx_abstract_model and flax.nnx.split to dynamically generate abstract parameters for NNX models when enable_nnx=True is passed in the config.
  2. Implemented parameter key restructuring to correctly map NNX trees (prepending params.params.) vs Linen trees (params.).
  3. Retained the legacy Linen Transformer initialization logic strictly as a fallback for older models (like Deepseekv4) that have not yet migrated to NNX.

Tests

1. Manually verified that passing `enable_nnx=True` successfully initializes the NNX abstract model and accurately resolves the params.params. key structure.

Command:

python3 src/maxtext/experimental/agent/ckpt_validation_pipeline/inspect_checkpoint.py maxtext base_config=src/maxtext/configs/models/deepseek4-tiny.yml run_name=test load_parameters_path=gs://dummy_path enable_nnx=True 
Total Parameters: 1104218112 (~1.10 B)

=== Structure ===
key: params.params.decoder.decoder_norm.scale | shape: (2048,)
key: params.params.decoder.layers.mlp.wi_0.kernel | shape: (2048, 16, 7168)
key: params.params.decoder.layers.mlp.wi_1.kernel | shape: (2048, 16, 7168)
key: params.params.decoder.layers.mlp.wo.kernel | shape: (7168, 16, 2048)
key: params.params.decoder.layers.post_self_attention_layer_norm.scale | shape: (2048, 16)
key: params.params.decoder.layers.pre_self_attention_layer_norm.scale | shape: (2048, 16)
key: params.params.decoder.layers.self_attention.key.kernel | shape: (2048, 16, 16, 128)
key: params.params.decoder.layers.self_attention.out.kernel | shape: (16, 16, 128, 2048)
key: params.params.decoder.layers.self_attention.query.kernel | shape: (2048, 16, 16, 128)
key: params.params.decoder.layers.self_attention.value.kernel | shape: (2048, 16, 16, 128)
key: params.params.decoder.logits_dense.kernel | shape: (2048, 32000)
key: params.params.token_embedder.embedding | shape: (32000, 2048)
INFO:absl:Peak Memory: 7.34 GB
2. Verified that falling back to older configs successfully routes to the legacy Linen Transformer initialization without failure.

Command:

python3 src/maxtext/experimental/agent/ckpt_validation_pipeline/inspect_checkpoint.py maxtext base_config=src/maxtext/configs/models/deepseek4-tiny.yml run_name=test load_parameters_path=gs://dummy_path enable_nnx=False
Total Parameters: 1104218112 (~1.10 B)

=== Structure ===
key: params.decoder.decoder_norm.scale | shape: (2048,)
key: params.decoder.layers.mlp.wi_0.kernel | shape: (2048, 16, 7168)
key: params.decoder.layers.mlp.wi_1.kernel | shape: (2048, 16, 7168)
key: params.decoder.layers.mlp.wo.kernel | shape: (7168, 16, 2048)
key: params.decoder.layers.post_self_attention_layer_norm.scale | shape: (2048, 16)
key: params.decoder.layers.pre_self_attention_layer_norm.scale | shape: (2048, 16)
key: params.decoder.layers.self_attention.key.kernel | shape: (2048, 16, 16, 128)
key: params.decoder.layers.self_attention.out.kernel | shape: (16, 16, 128, 2048)
key: params.decoder.layers.self_attention.query.kernel | shape: (2048, 16, 16, 128)
key: params.decoder.layers.self_attention.value.kernel | shape: (2048, 16, 16, 128)
key: params.decoder.logits_dense.kernel | shape: (2048, 32000)
key: params.token_embedder.embedding | shape: (32000, 2048)
INFO:absl:Peak Memory: 7.36 GB

Checklist

Before submitting this PR, please make sure (put X in square brackets):

  • I have performed a self-review of my code. For an optional AI review, add the gemini-review label.
  • I have necessary comments in my code, particularly in hard-to-understand areas.
  • I have run end-to-end tests tests and provided workload links above if applicable.
  • I have made or will make corresponding changes to the doc if needed, including adding new documentation pages to the relevant Table of Contents (toctree directive) as explained in our documentation.

@gemini-code-assist

gemini-code-assist Bot commented Jul 22, 2026

Copy link
Copy Markdown

Caution

The consumer version of Gemini Code Assist on GitHub has been sunset. All code review activity has officially ceased.

@olufiyin19
olufiyin19 force-pushed the feat/legacy-inspect-checkpoint-nnx-support branch from 0453b56 to c31d3ae Compare July 22, 2026 18:32
@codecov

codecov Bot commented Jul 22, 2026

Copy link
Copy Markdown

Codecov Report

✅ All modified and coverable lines are covered by tests.

📢 Thoughts on this report? Let us know!

shuningjin
shuningjin reviewed Jul 22, 2026
View reviewed changes

@shuningjin shuningjin left a comment •
edited
Loading

Copy link
Copy Markdown
Collaborator

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Thanks for adding the support for inspecting NNX checkpoint! By comparing the code with checkpoint_conversion/inspect_checkpoint.py, it seems the change is relatively safe (diff). Would you consider directly incorporating to that script for broad use?

@olufiyin19
olufiyin19 requested a review from parambole as a code owner July 22, 2026 23:36
@olufiyin19
olufiyin19 requested a review from shuningjin July 22, 2026 23:42
Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment

Reviewers

@abhinavclemson abhinavclemson Awaiting requested review from abhinavclemson abhinavclemson is a code owner
@gobbleturk gobbleturk Awaiting requested review from gobbleturk gobbleturk is a code owner
@khatwanimohit khatwanimohit Awaiting requested review from khatwanimohit khatwanimohit is a code owner
@bvandermoon bvandermoon Awaiting requested review from bvandermoon bvandermoon is a code owner
@vipannalla vipannalla Awaiting requested review from vipannalla vipannalla is a code owner
@RissyRan RissyRan Awaiting requested review from RissyRan RissyRan is a code owner
@richjames0 richjames0 Awaiting requested review from richjames0 richjames0 is a code owner
@gagika gagika Awaiting requested review from gagika gagika is a code owner
@shralex shralex Awaiting requested review from shralex shralex is a code owner
@SurbhiJainUSC SurbhiJainUSC Awaiting requested review from SurbhiJainUSC SurbhiJainUSC is a code owner
@hengtaoguo hengtaoguo Awaiting requested review from hengtaoguo hengtaoguo is a code owner
@A9isha A9isha Awaiting requested review from A9isha A9isha is a code owner
@aireenmei aireenmei Awaiting requested review from aireenmei aireenmei is a code owner
@NuojCheng NuojCheng Awaiting requested review from NuojCheng NuojCheng is a code owner
@jiangjy1982 jiangjy1982 Awaiting requested review from jiangjy1982 jiangjy1982 is a code owner
@suexu1025 suexu1025 Awaiting requested review from suexu1025 suexu1025 is a code owner
@dipannita08 dipannita08 Awaiting requested review from dipannita08 dipannita08 is a code owner
@igorts-git igorts-git Awaiting requested review from igorts-git igorts-git is a code owner
@darisoy darisoy Awaiting requested review from darisoy darisoy is a code owner
@xibinliu xibinliu Awaiting requested review from xibinliu xibinliu is a code owner
@huytransformer huytransformer Awaiting requested review from huytransformer huytransformer is a code owner
@entrpn entrpn Awaiting requested review from entrpn
@Rohan-Bierneni Rohan-Bierneni Awaiting requested review from Rohan-Bierneni
@parambole parambole Awaiting requested review from parambole parambole is a code owner
@shuningjin shuningjin Awaiting requested review from shuningjin shuningjin is a code owner

At least 2 approving reviews are required to merge this pull request.

Assignees

No one assigned

Labels

None yet

Projects

None yet

Milestone

No milestone

Development

Successfully merging this pull request may close these issues.

2 participants

@olufiyin19 @shuningjin